# -*- coding: utf-8 -*-
# nlp_client.py
# Created by Hardy on 23th, Jan
# Copyright 2017 杭州网川教育科技有限公司. All rights reserved.

import requests
from utils.utils import load_pickle

KEYWORDCLOUD_TPL = 'http://%s:%s/keywordcloud'
KEYWORD_TPL = 'http://%s:%s/keyword'
WORDCLOUD_TPL = 'http://%s:%s/wordcloud'
WORDSEG_TPL = 'http://%s:%s/seg'
CLASSIFY_TPL = 'http://%s:%s/classify'
GET_WORD_VECTORS_TPL = 'http://%s:%s/word_vectors'

ARG_TEXT = 'text'
ARG_KEYWORDS = 'keywords'
ARG_TOP_K = 'keyword_top_k'
ARG_TOP_N = 'keyword_cloud_n'
ARG_ALLOW_POS = 'keyword_allow_pos'

RESULT = 'result'
MESSAGE = 'message'


class NLPService(object):

    def __init__(self, host, port, timeout=1, defaults=None, brand_pickle=None):
        self._keywordcloud_url = KEYWORDCLOUD_TPL % (host, port)
        self._keyword_url = KEYWORD_TPL % (host, port)
        self._classify_url = CLASSIFY_TPL % (host, port)
        self._wordseg_url = WORDSEG_TPL % (host, port)
        self._wordcloud_url = WORDCLOUD_TPL % (host, port)
        self._get_word_vectors_url = GET_WORD_VECTORS_TPL % (host, port)
        self._timeout = timeout
        self._defaults = defaults
        self.brand_kv = load_pickle(brand_pickle) if brand_pickle is not None else {}
        if self._defaults is None:
            self._defaults = {}
        if self._defaults.get('keyword_top_k', None) is None:
            self._defaults['keyword_top_k'] = None
        if self._defaults.get('keyword_cloud_n', None) is None:
            self._defaults['keyword_cloud_n'] = None
        if self._defaults.get('keyword_allow_pos', None) is None:
            self._defaults['keyword_allow_pos'] = None

    def keywords(self, text, keyword_top_k=None, keyword_allow_pos=None):
        keyword_top_k, keyword_allow_pos, wordcloud_topn = self._parseDefault(keyword_top_k=keyword_top_k, keyword_allow_pos=keyword_allow_pos)
        res = requests.post(self._keyword_url, json={ARG_TEXT: text, ARG_TOP_K: keyword_top_k, ARG_ALLOW_POS: keyword_allow_pos}, timeout=self._timeout)
        res = res.json()
        if RESULT in res:
            return res[RESULT]
        else:
            raise RuntimeError(res.json().get(MESSAGE, 'no message found'))

    def seg(self, text,  keyword_allow_pos=None):
        keyword_top_k, keyword_allow_pos, wordcloud_topn = self._parseDefault(keyword_allow_pos=keyword_allow_pos)
        res = requests.post(self._keyword_url, json={ARG_TEXT: text, ARG_ALLOW_POS: keyword_allow_pos}, timeout=self._timeout)
        res = res.json()
        if RESULT in res:
            return res[RESULT]
        else:
            raise RuntimeError(res.json().get(MESSAGE, 'no message found'))

    def get_brands(self, keywords):
        return [x for x in keywords if self.brand_kv.get(x) is not None]

    def classify(self, keyword_list):
        res = requests.post(self._classify_url, json={ARG_KEYWORDS: keyword_list}, timeout=self._timeout)
        res = res.json()
        if RESULT in res:
            return res[RESULT]
        else:
            raise RuntimeError(res.json().get(MESSAGE, 'no message found'))

    def keywordcloud(self, text, keyword_top_k=None, keyword_allow_pos=None, wordcloud_topn=None):
        keyword_top_k, keyword_allow_pos, wordcloud_topn = self._parseDefault(keyword_top_k, keyword_allow_pos, wordcloud_topn)
        res = requests.post(self._keywordcloud_url, json={ARG_TEXT: text, ARG_TOP_K: keyword_top_k, ARG_ALLOW_POS: keyword_allow_pos, ARG_TOP_N: wordcloud_topn}, timeout=self._timeout)
        res = res.json()
        if RESULT in res:
            return res[RESULT]
        else:
            raise RuntimeError(res.json().get(MESSAGE, 'no message found'))

    def wordcloud(self, text, keyword_top_k=None, keyword_allow_pos=None, wordcloud_topn=None):
        keyword_top_k, keyword_allow_pos, wordcloud_topn = self._parseDefault(keyword_cloud_n=wordcloud_topn)
        res = requests.post(self._wordcloud_url, json={ARG_TEXT: text, ARG_TOP_N: wordcloud_topn}, timeout=self._timeout)
        res = res.json()
        if RESULT in res:
            return res[RESULT]
        else:
            raise RuntimeError(res.json().get(MESSAGE, 'no message found'))

    def get_word_vectors(self, texts):
        res = requests.post(self._get_word_vectors_url, json={'arg_dict': {'texts': texts}}, timeout=self._timeout)
        res = res.json()
        if RESULT in res:
            return res[RESULT]['word2vec']
        else:
            raise RuntimeError(res.json().get(MESSAGE, 'no message found'))

    def _parseDefault(self, keyword_top_k=None, keyword_allow_pos=None, keyword_cloud_n=None):
        if keyword_top_k is None:
            keyword_top_k = self._defaults['keyword_top_k']
        if keyword_allow_pos is None:
            keyword_allow_pos = self._defaults['keyword_allow_pos']
        if keyword_cloud_n is None:
            keyword_cloud_n = self._defaults['keyword_cloud_n']
        return int(keyword_top_k), keyword_allow_pos, int(keyword_cloud_n)